import logging
import json
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
import soundfile
import os

# 配置日志
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler())

os.environ["SM_MODEL_DIR"] = "/root/.cache/modelscope/hub/models/"
model_path = os.environ["SM_MODEL_DIR"]
logger.info("Libraries are loaded")

def get_device():
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    return device

def load_model():
    try:
        device = get_device()
        model = Wav2Vec2ForCTC.from_pretrained(model_path).to(device)
        logger.info("Model is loaded")
        return model
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        return None

def process_input(json_request_data):
    try:
        input_data = json.loads(json_request_data)
        logger.info("Input data is processed")
        return input_data
    except json.JSONDecodeError as e:
        logger.error(f"Failed to decode input JSON: {e}")
        return None

def perform_prediction(input_data, model):
    if input_data is None or model is None:
        return None
    try:
        logger.info("Starting inference.")
        device = get_device()
        #logger.info(input_data)
        speech_array = input_data['speech_array']
        sampling_rate = input_data['sampling_rate']
        processor = Wav2Vec2Processor.from_pretrained(model_path)
        input_values = processor(speech_array, sampling_rate=sampling_rate, return_tensors="pt").input_values.to(device)
        with torch.no_grad():
            logits = model(input_values).logits
        pred_ids = torch.argmax(logits, dim=-1)
        transcript = processor.batch_decode(pred_ids)[0]
        return transcript
    except Exception as e:
        logger.error(f"Prediction failed: {e}")
        return None

def prepare_output(transcript):
    if transcript is None:
        return None
    try:
        return json.dumps(transcript), 'application/json'
    except Exception as e:
        logger.error(f"Failed to prepare output JSON: {e}")
        return None

def predict(json_request_data):
    model = load_model()
    input_data = process_input(json_request_data)
    transcript = perform_prediction(input_data, model)
    output = prepare_output(transcript)
    return output

# 模拟输入的 JSON 数据
file_name ='MKH800_19_0001.wav'

speech_array, sampling_rate = soundfile.read(file_name)
json_request_data = {"speech_array": speech_array.tolist(),
                     "sampling_rate": sampling_rate}
# 将字典转换为 JSON 字符串
json_request_str = json.dumps(json_request_data)
result = predict(json_request_str)
if result:
    print(result[0])
